import os
from pathlib import Path
import json
import pickle
import networkx as nx
import matplotlib.pyplot as plt
import cv2
import open3d as o3d
from semantic_nav.slam_utils import pose_dist, pose_minus, pose_np2dict, numpy_to_o3d_pcd
from semantic_nav.log_utils import get_logger
logger = get_logger()


def show_graph(graph):
    pos = nx.spring_layout(graph)  # positions for all nodes
    nx.draw(graph, pos, with_labels=True, font_weight='bold')
    plt.show()


def show_graph_with_images(graph, images):
    fig, ax = plt.subplots(figsize=(20, 20))  # Specify image display size (width x height)
    layout_scale = 100.0
    image_size_factor = 0.2
    image_size = image_size_factor * layout_scale
    pos = nx.shell_layout(graph, scale=layout_scale)  # positions for all nodes
    for idx, loc in pos.items():
        # `extent=[left, right, bottom, top]`
        ax.imshow(images[idx], extent=[loc[0]-image_size, loc[0]+image_size, loc[1]-image_size, loc[1]+image_size])

    nx.draw(graph, pos=pos, ax=ax, with_labels=True, font_weight='bold')
    plt.show()


class TopoMap:
    def __init__(self, path, debug: bool=False):
        self.map_dir = path
        self.graph_path = os.path.join(path, 'graph.pkl')
        self.observation_dir = os.path.join(path, 'observations')
        self.observation_dict = {}
        Path(self.observation_dir).mkdir(parents=True, exist_ok=True)
        if os.path.exists(self.graph_path):
            with open(self.graph_path, 'rb') as f:
                self.graph = pickle.load(f)
            self.node_num = self.graph.number_of_nodes()
            self.last_pose = self.graph.nodes[self.node_num - 1]['pose'] if self.node_num > 0 else None
            frame_dirs = [p for p in Path(self.observation_dir).iterdir() if p.is_dir()]
            self.frame_num = len(frame_dirs)
        else:
            self.graph = nx.Graph()
            self.last_pose = None
            self.node_num = 0
            self.frame_num = 0
        self.debug = debug
        if self.debug:
            self.show_with_images()

    def __del__(self):
        with open(self.graph_path, 'wb') as f:
            pickle.dump(self.graph, f, protocol=pickle.HIGHEST_PROTOCOL)
        for node in self.graph.nodes:
            self.graph.nodes[node]['pose'] = pose_np2dict(*self.graph.nodes[node]['pose'])
        for edge in self.graph.edges:
            self.graph.edges[edge]['relative_pose'] = pose_np2dict(*self.graph.edges[edge]['relative_pose'])
        json_data = nx.readwrite.json_graph.node_link_data(self.graph)
        with open(os.path.join(self.map_dir, 'graph.json'), 'w') as f:
            json.dump(json_data, f)
        if self.debug:
            self.show_with_images()

    def add_node(self, node, **attr):
        self.graph.add_node(node, **attr)

    def add_edge(self, u, v, **attr):
        self.graph.add_edge(u, v, **attr)

    def has_node(self, node):
        return self.graph.has_node(node)

    def has_edge(self, u, v):
        return self.graph.has_edge(u, v)

    def node(self, node_idx):
        return self.graph.nodes[node_idx]

    def nodes(self, data=False):
        return self.graph.nodes(data=data)

    def number_of_nodes(self):
        return self.graph.number_of_nodes()

    def edge(self, u, v):
        return self.graph.edges[u, v]

    def edges(self, data=False):
        return self.graph.edges(data=data)

    def number_of_edges(self):
        return self.graph.number_of_edges()

    def grow_graph(self, this_node, last_node, this_pose):
        if this_node is None:
            return
        if not self.has_node(this_node):
            self.add_node(this_node, pose=this_pose, frame_idx=self.frame_num - 1)

        if self.last_pose is not None and last_node is not None:
            rot_dist, trans_dist = pose_dist(this_pose, self.last_pose)
            dist = 2.0 * abs(rot_dist) + abs(trans_dist)
            relative_pose = pose_minus(this_pose, self.last_pose)
            self.add_edge(last_node, this_node, relative_pose=relative_pose, dist=dist)

        self.last_pose = this_pose

    def save_observation(self, observations):
        this_frame_dir = os.path.join(self.observation_dir, f'{self.frame_num}')
        self.frame_num += 1
        Path(this_frame_dir).mkdir(parents=True, exist_ok=True)
        for i, (cv_color_image, cv_depth_image, point_cloud) in enumerate(observations):
            cv2.imwrite(os.path.join(this_frame_dir, f"color_{i}.png"), cv_color_image)
            cv2.imwrite(os.path.join(this_frame_dir, f"depth_{i}.png"), cv_depth_image)
            o3d.io.write_point_cloud(os.path.join(this_frame_dir, f"depth_{i}.ply"), point_cloud, write_ascii=True)

    def load_observations(self):
        for node in self.graph.nodes():
            self.get_frame(self.graph.nodes[node]['frame_idx'])
        return self.observation_dict

    def get_frame(self, frame_idx):
        frame_dir = os.path.join(self.observation_dir, f'{frame_idx}')
        if frame_idx in self.observation_dict:
            return self.observation_dict[frame_idx]
        img_names = [os.path.join(frame_dir, filename) for filename in os.listdir(frame_dir) if filename.startswith('color')]
        img_names.sort()
        images = [cv2.imread(img_name, cv2.IMREAD_UNCHANGED) for img_name in img_names]

        depth_names = [os.path.join(frame_dir, filename) for filename in os.listdir(frame_dir) if filename.startswith('depth') and filename.endswith('.png')]
        depth_names.sort()
        depth_images = [cv2.imread(img_name, cv2.IMREAD_UNCHANGED) for img_name in depth_names]

        points_names = [os.path.join(frame_dir, filename) for filename in os.listdir(frame_dir) if filename.startswith('depth') and filename.endswith('.ply')]
        points_names.sort()
        points_array = [o3d.io.read_point_cloud(points_name) for points_name in points_names]

        frame_observation = list(zip(images, depth_images, points_array))
        self.observation_dict[frame_idx] = frame_observation

        return frame_observation

    def find_path(self, start_idx, end_idx, weight='dist'):
        try:
            path = nx.shortest_path(self.graph, source=start_idx, target=end_idx, weight=weight)
        except nx.NetworkXNoPath:
            logger.error("No path found")
            return None
        return path

    def show(self):
        show_graph(self.graph)

    def show_with_images(self):
        self.load_observations()
        images = {}
        for node in self.graph.nodes():
            frame_idx = self.graph.nodes[node]['frame_idx']
            frame_observation = self.get_frame(frame_idx)
            images[node] = cv2.cvtColor(frame_observation[0][0], cv2.COLOR_BGR2RGB)

        show_graph_with_images(self.graph, images)
